from django.core.management.base import BaseCommand
from django.contrib.auth.models import Group, Permission
from django.contrib.contenttypes.models import ContentType
from monitor.models import Asset, Alert, Device, IDC, NetworkDevice, Host, Rack, Connection

class Command(BaseCommand):
    help = '创建用户组和设置权限'

    def handle(self, *args, **options):
        # 创建用户组
        admin_group, _ = Group.objects.get_or_create(name='管理员')
        normal_group, _ = Group.objects.get_or_create(name='普通用户')
        readonly_group, _ = Group.objects.get_or_create(name='只读用户')

        # 获取所有模型的权限
        models = [Asset, Alert, Device, IDC, NetworkDevice, Host, Rack, Connection]
        
        # 管理员组权限（所有权限）
        admin_permissions = []
        for model in models:
            content_type = ContentType.objects.get_for_model(model)
            permissions = Permission.objects.filter(content_type=content_type)
            admin_permissions.extend(permissions)
        
        # 普通用户组权限（查看所有，修改部分）
        normal_permissions = []
        for model in models:
            content_type = ContentType.objects.get_for_model(model)
            # 添加查看权限
            view_perm = Permission.objects.get(
                content_type=content_type,
                codename=f'view_{model._meta.model_name}'
            )
            normal_permissions.append(view_perm)
            
            # 对于告警和设备，添加修改权限
            if model in [Alert, Device, Host, NetworkDevice]:
                change_perm = Permission.objects.get(
                    content_type=content_type,
                    codename=f'change_{model._meta.model_name}'
                )
                normal_permissions.append(change_perm)
        
        # 只读用户组权限（只有查看权限）
        readonly_permissions = []
        for model in models:
            content_type = ContentType.objects.get_for_model(model)
            view_perm = Permission.objects.get(
                content_type=content_type,
                codename=f'view_{model._meta.model_name}'
            )
            readonly_permissions.append(view_perm)

        # 设置权限
        admin_group.permissions.set(admin_permissions)
        normal_group.permissions.set(normal_permissions)
        readonly_group.permissions.set(readonly_permissions)

        # 添加用户管理相关权限给管理员组
        user_ct = ContentType.objects.get(app_label='auth', model='user')
        group_ct = ContentType.objects.get(app_label='auth', model='group')
        
        admin_group.permissions.add(
            *Permission.objects.filter(content_type__in=[user_ct, group_ct])
        )

        self.stdout.write(self.style.SUCCESS('成功创建用户组和设置权限'))

        # 输出每个组的权限数量
        self.stdout.write(f"管理员组权限数量: {admin_group.permissions.count()}")
        self.stdout.write(f"普通用户组权限数量: {normal_group.permissions.count()}")
        self.stdout.write(f"只读用户组权限数量: {readonly_group.permissions.count()}") 